-------------------------------------------------------------------------------
-- Copas2 - Coroutine Oriented Portable Asynchronous Services, take 2
-- 
-- This is a rewrite of the original COPAS but using libevent as the
-- signalling mechanism.
--
-- The async IO wrapper(alas copas.send/receive) has also been taken out
-- to a seperate module for case where people can use other dispatcher other
-- than COPAS. It supports other io type like pipe and serialized read/write
-- among threads(for HTTP pipelining and other duplex protocols
-- 
-- Offers a dispatcher for non-blocking io based on coroutines and libevent.
-- Usage:
--    copas.addserver(server, handler, timeout)
--    copas.addthread(thread, ...) Create a new coroutine thread and run it with args
--    copas.loop(timeout, once) - listens infinetely or once
--    copas.step(timeout) - executes one listening step(COPAS compatibility)
--    copas.wrap  - wraps an IO object to co-op with the dispatcher
--    copas.connect - blocks only the thread until connection completes
--
--
-- Author : Gary Ng<linux@garyng.com>
-- 
-- Portion of the code is taken from COPAS with the following authors.
--
-- Authors: Andre Carregal and Javier Guerra
-- Contributors: Diego Nehab, Mike Pall, David Burgess, Leonardo Godinho
--               and Thomas Harning Jr.
--
--
-- Copyright 2007 - Gary Ng<linux@garyng.com>
-- This code can be distributed under LGPL license
--
-------------------------------------------------------------------------------
local socket = require "socket"
local asyncio = require 'asyncio'
local lasync = require 'lasync'
require'coxpcall'
local time = os.time
local duration = os.difftime
local libevent = lasync.handle()
local cocreate,coresume = coroutine.create, coroutine.resume
local tinsert, tremove = table.insert, table.remove
local type=type
local unpack = unpack
local pairs,ipairs = pairs,ipairs
asyncio.CHECK_STALLED = false

WATCH_DOG_TIMEOUT = 120
FAST_ACCEPT = false
RECREATE_EVENT = true

-- Following section is taken from COPAS
-- Begin COPAS
-- Redefines LuaSocket functions with coroutine safe versions
-- (this allows the use of socket.http from within copas)
local function statusHandler(status, ...)
  if status then return ... end
  return nil, ...
end

function socket.protect(func)
  return function (...)
		local ret = {copcall(func, ...)}
		local status = table.remove(ret,1)
		if status then
			return unpack(ret)
		end
		return nil, unpack(ret)
	end
end

function socket.newtry(finalizer)
	return function (...)
    local arg={...}
		local status = arg[1]or false
		if (status==false)then
			table.remove(arg,1)
			local ret = {pcall(finalizer, unpack(arg) ) }
			error(arg[1], 0) --this would terminate a running thread(coroutine)
		end
		return unpack(arg)
	end
end
-- end of LuaSocket redefinitions

-- End COPAS
module ("copas2", package.seeall)


local weak_key = {__mode = "k"}
local weak_val = {__mode = "v"}
local weak_table = {__mode = "kv"}
local _servers = {} -- listening sockets
--local _threads = setmetatable({}, weak_key) -- pending threads indexed by fd
local _threads = {} -- pending threads indexed by event object
local _dangling = {}
local _killing = {}


local fd_map=setmetatable({},{__index="v"})

function state()
  return libevent:state()
end

function wrap (skt, buffered_write)
	local s =  asyncio.wrap(skt, buffered_write)
  --s.event = libevent:event(skt:getfd(), skt):open()
  local old_close = s.close
  s.close = function(wrapped_skt) 
            if wrapped_skt.event then wrapped_skt.event:close() end
            if old_close then old_close(s) end
            end
  return s
end

--an ojbect mapping if the send/receive method is used
--instead of the wrapper

local obj_map=setmetatable({},{__mode="k"}) 

function send(obj, data)
    local obj_key = obj
    local state = obj_map[obj_key]
    if not state then 
        state = wrap(obj)
        obj_map[obj_key] = state
    end

    return state:send(data)
end

function receive(obj, pattern, ...)
    local obj_key = obj
    local state = obj_map[obj_key]
    if not state then 
        state = wrap(obj)
        obj_map[obj_key] = state
    end

    return state:receive(pattern, ...)
end

-- waits until connection is completed
function connect(skt, host, port)
  if non_block_connect == false then 
    skt:settimeout()
    return skt:connect(host,port)
  end
  local obj_key = skt
  local state = obj_map[obj_key]
  if not state then 
      state = wrap(skt)
      obj_map[obj_key] = state
  end
  return state:connect(host, port)
end

local timer_mt={
    sleep = function(self, time, conn)
        self.host = conn
        self.event:open(time)
        coroutine.yield(self,"sleeping")  
    end,
    start = function(self, time)
        self.event:open(time)
    end,
    close = function(self, ...)
        self.event:close()
    end,
    getfd = function(self, ...)
        return self
    end,
    }; timer_mt.__index = timer_mt 

function timer(once)
    local timer = setmetatable({}, timer_mt)
    local ev = libevent:event(-1, timer, once and 1 or 0)
    timer.event = ev
    return timer
end

function sleep(timeout,once)
    local t = timer(once)
    t:sleep(timeout)
    if once then t:close() end
end

--------------------------------------------------
-- Error handling
--------------------------------------------------

local _errhandlers = {}   -- error handler per coroutine

function setErrorHandler (err)
	local co = coroutine.running()
	if co then
		_errhandlers [co] = err
	end
end

local function _deferror (msg, co, skt)
	--print (msg, co, skt)
end

-------------------------------------------------------------------------------
-- Thread handling
-------------------------------------------------------------------------------

local function _doTick (new, co, skt, ...)
	if not co then return end
	
	local ok, res, q = coresume(co, skt, ...)

	if not ok or not res or _killing[co] then
        local killed = _killing[co]
        if co then 
          _dangling[co] = nil 
          _killing[co] = nil 
          if coroutine.status(co) ~= "dead" then
            return _doTick(false, co, res, true)
          end
        end
		    if not ok then pcall (_errhandlers [co] or _deferror, res, co, skt) end
		    _errhandlers [co] = nil

        if skt then
            local ev = skt.event
            if ev then
                local pending = _threads[ev]
                if true or pending then
                    if #pending.reading + #pending.writing + #pending.sleeping > 0 then
                        return
                    end
                    _threads[ev] = nil 
                end
                if (skt.getfd and type(skt:getfd()) == "number") then
                    skt.event:close() 
                    skt.event = nil
                end
            end
            --skt:close()
        end
  else
    if skt and skt.event then
      local pending = _threads[skt.event]
      if pending then
        local active_threads = #pending.reading + #pending.writing + #pending.sleeping
        if active_threads < 1 then 
          _threads[skt.event] = nil 
          if (skt.getfd and type(skt:getfd()) == "number") then
              skt.event:close() 
              skt.event = nil
          end
        end
      end
    end
    local dead = co and _dangling[co]
    if not dead or os.time() - dead < WATCH_DOG_TIMEOUT*2 then
        local ev = res.event
        if not ev then
            --if skt and skt:getfd() == res:getfd() then ev = skt.event end
            ev = ev or libevent:event(res:getfd(), res):open()
            res.event = ev
        end
        local pending = _threads[ev]
        if not pending then
            pending = {reading={},writing={},sleeping={}}
           _threads[ev] = pending
        end
        if new then 
            tinsert(pending[q], {co, res, time()})
        else 
            tinsert(pending[q],1, {co, res, time()}) 
        end
    else
      --_killing[co] = os.time() 
      --_dangling[co] = nil
      --res:close()
      return _doTick(false, co, res, true)
    end
	end
end

-- accepts a connection on socket input
local function _accept(input, handler)
	local client
  if input.receivefrom then client = input
  else client = input:accept() end
	if client then 
		client:settimeout(0)
		local co = cocreate(handler[1])
	    return co, client
	end
end

local function _wakeup(q, i)
    local job = tremove(q,i)
    if job then
        local co,skt,t = unpack(job) 
        _doTick(false, co, skt)
        return true
    end
end

-------------------------------------------------------------------------------
-- Adds a server/handler pair to Copas dispatcher
-------------------------------------------------------------------------------
function addserver(server, handler, timeout)
    local fd = server:getfd()
    local ev = libevent:event(fd, server):open()

    server:settimeout(timeout or 0)
    _servers[server] = {handler, ev}
end

function delserver(server)
    local _,ev = unpack(_servers[server] or {})
    if ev then ev:close() end
    _servers[server] = nil
end

-------------------------------------------------------------------------------
-- Adds an new courotine thread to Copas dispatcher
-------------------------------------------------------------------------------
function addthread(thread, ...)
	local co = cocreate(thread)
	_doTick (true, co, nil, ...)
    return co
end

-------------------------------------------------------------------------------
-- The main loop
-------------------------------------------------------------------------------
function cleansing()
  local active_events = 0
  local reading = 0
  local writing = 0
  local sleeping = 0
  local now = os.time()
  local t={}
  for k, v in pairs(_threads) do t[k] = v end --must copy them out
  for k, pending in pairs(t) do
      active_events = active_events + 1
      local suspicious = 0
      for i, v in ipairs(pending.reading) do
          reading = reading + 1
          local co, skt = unpack(v)
          if type(skt) ~= "userdata" then
              if not skt.dangling then skt.dangling = now
              elseif os.difftime(now, skt.dangling) > WATCH_DOG_TIMEOUT then
                  if not _dangling[co] then _dangling[co] = now end
                  suspicious = suspicious + 1
                  --local c,e = lasync.probe(skt:getfd())
                  --if e then _wakeup(pending.reading, 1) end
                  _wakeup(pending.reading,1)
              end
          end
      end
      --if suspicious > 0 then print("suspcious read", suspicious) end
      local suspicious = 0
      for i, v in ipairs(pending.writing) do
          writing = writing + 1
          local co, skt = unpack(v)
          if not skt.dangling then skt.dangling = now
          elseif os.difftime(now, skt.dangling) > WATCH_DOG_TIMEOUT then
              if not _dangling[co] then _dangling[co] = now end
              suspicious = suspicious + 1
              --local c,e = lasync.probe(skt:getfd())
              --if e then _wakeup(pending.writing, 1) end
              _wakeup(pending.writing,1)
          end
      end
      --if suspicious > 0 then print("suspcious write", suspicious) end
      local suspicious = 0
      for i, v in ipairs(pending.sleeping) do
          sleeping = sleeping + 1
          local co, skt = unpack(v)
          if not skt.dangling then skt.dangling = now
          elseif os.difftime(now, skt.dangling) > WATCH_DOG_TIMEOUT then
              suspicious = suspicious + 1
              --local c,e = lasync.probe(skt:getfd())
              --if e then _wakeup(pending.writing, 1) end
          end
      end
      --if suspicious > 0 then print("suspcious sleep", suspicious) end
  end
  if active_events > 0 then
      --print("active events", active_events, reading, writing, sleeping)
  end
end

local watchdog = timer(false)
function loop(timeout, once)
    local l={}
    local sselect = socket.select

    for k,v in pairs(_servers) do 
        l[#l+1] = k
    end

    local function check_accept(skt)
        local r,w

        if skt then r={skt}
        else r,w=sselect(l,{},0) end

        for k,v in ipairs(r) do 
            local co, client = _accept(v, _servers[v])
            if co then
		        _doTick (true, co, client)
            end
        end
    end

    local dispatcher = {
        ["readable"] = function(self, ev)
            local skt = ev:object()
            if skt then
                if _servers[skt] then
                    check_accept(skt)
                else
                    local pending = _threads[ev]
                    if pending and not _wakeup(pending.reading, 1) and #pending.writing > 0 then
                        local co, v, t = unpack(pending.writing[1])
                        if duration(time(), t) > 120 then
                            --print(time(), "readable but in waiting for write queue", duration(time(), t))
                            --lasync.probe(skt:getfd())
                        end
                    end
                end
            end
        end,
        ["writable"] = function(self, ev)
            local skt = ev:object()
            if skt then
                local pending = _threads[ev]
                if pending and not _wakeup(pending.writing, 1) and #pending.reading > 0 then
                    local co, v, t = unpack(pending.reading[1])
                    if duration(time(), t) > 120 then
                        --print(time(), "writiable but in waiting for read queue", duration(time(), t))
                        --lasync.probe(skt:getfd())
                        _wakeup(pending.reading, 1)
                    end
                end
            end
            if FAST_ACCEPT then check_accept() end
        end,
        ["timeout"] = function(self, ev)
            local skt = ev:object()
            if skt then
                local pending = _threads[ev]
                if pending then
                    _wakeup(pending.sleeping, 1)
                    if type(skt:getfd()) ~= "number" then
                        _threads[ev] = nil
                        skt:close()
                    end
                end
            end
            if skt == watchdog then
              cleansing()
            end
        end,
        }
    watchdog:start(WATCH_DOG_TIMEOUT)
    local r = libevent:loop(dispatcher,once and 1 or 0)
    cleansing()
    watchdog:close()
    return r
end

function release(skt)
  if skt then
    local wrapper = obj_map[skt]
    if wrapper and wrapper.event then wrapper.event:close() end
    if skt.event then skt.event:close() end
  end
end

function loopexit()
 libevent:loopexit() 
end

function step(timeout)
    return loop(timeout, true)
end

